module efield

  use dimensions
  use fft_3d

  complex, private   :: phi03DFT(jz,jy,jx)  ! FT of potential of background Efield;
  real,    private   :: kk3D(jz,jy,jx)      ! matrix of k-squared

  real,    private   :: avEx(jy,jx)         ! z-averaged spectrum, diagnostics
  real,    private   :: avEy(jy,jx)         ! z-averaged spectrum, diagnostics
  real,    private   :: avEz(jy,jx)         ! z-averaged spectrum, diagnostics

  complex, private   :: densityFTzC(jz)     ! FT in z at the center (x=0), diagnostics
  complex, private   :: E0C(jz)             ! background E-field at the center
  complex, private   :: E0zC(jz)            ! = E0C up to center definition
  complex, private   :: Phi0T               ! complex amplitude of background potential

  complex, private   :: xMod(jx)            ! modulation vectors 
  complex, private   :: yMod(jy)


contains

!--------------------------------------------------------------------------------------
subroutine efield_init()

! initialized private "phi03DFT", "E0C", "E0zC", and "Phi0T"
!
!--------------------------------------------------------------------------------------

  use fft_wrap, only   : fft_forward, fft_backward
  use parameters, only : phi0, kz, iE0Modulation, SpeckleWidth, SpeckleDepth, iunit

  implicit none

  integer              :: ix, iy, iz

  complex              :: Phi1Dx(jx)
  complex              :: Phi1Dz(jz)

      !-- initialization of background field 

      phi0T = phi0

      do iz=1,jz
         E0C(iz) =  -iunit * kz * phi0 *exp ( iunit * kz * (iz-1) * dz  )
      end do

      call fft_forward(jz,E0C,E0C)
      E0C=E0C/float(jz)
      call dealias(jz,E0C)
      call fft_backward(jz,E0C,E0C)



      !-- setting modulation vectors --

      xMod(:) = exp( - 20 * xvec(:)**2 / SpeckleWidth**2 )
      yMod(:) = exp( - 20 * yvec(:)**2 / SpeckleDepth**2 )

      call fft_forward(jx,xMod,xMod)
      xMod=xMod/float(jx)
      call deAlias(jx,xMod)
      call fft_backward(jx,xMod,xMod)  

      call fft_forward(jy,yMod,yMod)
      yMod=yMod/float(jy)
      call deAlias(jy,yMod)
      call fft_backward(jy,yMod,yMod)  


      !-- modulate source --
      !
      !  modulate INITIAL source in x direction
      !  E0C(iz) =  -iunit * kz * phi0 *exp ( iunit * kz * (iz-1) * dz  )     
      !  so when x=SpeckleWidth/2, E0z is down by a factor of 0.00674
      !  E0C(iz) = -iunit * kz * phi0 *exp ( iunit * kz * z  )

      if (iE0Modulation) then
         do ix=1,jx
            do iy=1,jy
               phi03DFT(:,iy,ix) = xMod(ix)*yMod(iy)*phi0 *exp ( iunit * kz * zvec(:)  ) 
            enddo
        enddo
        E0zC = xMod(jx/2+1)*yMod(jy/2+1)*E0C
      else
         do ix=1,jx
            do iy=1,jy
               phi03DFT(:,iy,ix) = phi0 *exp ( iunit * kz * zvec(:)  )
            end do
         end do
         E0zC  = E0C
      endif  


      !-- take Fourier transform of initial potential, store result

      call fft_forward_x(phi03DFT)
      call fft_forward_y(phi03DFT)
      call fft_forward_z(phi03DFT)

      !-- generate matrix of k-squared

      do ix=1,jx
         do iy=1,jy
            do iz=1,jz
               kk3D(iz,iy,ix) = kXvec(ix)**2 + kZvec(iz)**2  + kYvec(iy)**2
            end do
         end do
      end do

end subroutine efield_init


!-------------------------------------------------------------------------------------
subroutine efield_compute_E(dfu, ExFilter, EyFilter, &            ! inputs 
                            Ez3D, Ex3D, Ey3D, OmegaBounce)        ! outputs
!-------------------------------------------------------------------------------------

  use fft_wrap,   only : fft_forward, fft_backward
  use parameters, only : iunit, nz

  implicit none


  !-- external variables --

  real,    intent(in)  ::  dfu(jv,jz,jy,jx,jDirection) 
  real,    intent(in)  ::  ExFilter(jx)
  real,    intent(in)  ::  EyFilter(jy)

  real,    intent(out) ::  Ez3D(jz,jy,jx)
  real,    intent(out) ::  Ex3D(jz,jy,jx)
  real,    intent(out) ::  Ey3D(jz,jy,jx)
  real,    intent(out) ::  OmegaBounce


  !-- internal variables --

  complex  phi3DFT(jz,jy,jx)
  complex  dummy3D(jz,jy,jx)
  complex  dummy1Dz(jz)

  real     PhiCenter
  integer  ix, iy, iz, iv, iDirection

      
  !-- integrate density distribution, store as dummy3D --

  do ix=1,jx
     do iy=1,jy
        do iz=1,jz
           dummy3D(iz,iy,ix) = sum(dfu(:,iz,iy,ix,:)*dv) 
        end do
     end do
  end do


  !-- find FT of potential --
 
  call fft_forward_z(dummy3D)

  dummy3D(1,:,:) = 0.0
        
  densityFTzC(:) = dummy3D(:,jy/2+1,jx/2+1)       ! save for diagnostic

  call fft_forward_x(dummy3D)
  call fft_forward_y(dummy3D)

  do ix=1,jx
     dummy3D(:,:,ix) = dummy3D(:,:,ix)*ExFilter(ix)
  enddo

  do iy=1,jy
     dummy3D(:,iy,:) = dummy3D(:,iy,:)*EyFilter(iy)
  enddo


  phi3DFT = dummy3D / kk3D                        ! inversing Laplacian in k-space
 
  phi3DFT = phi3DFT + phi03DFT                    ! total potential

  phi3DFT(:,1,1) = 0                              ! spatially uniform in x not allowed (03/21/2011)
        
 

  !-- differentiate potential to find z-component of electric field -- 

  do iz=1,jz
     dummy3D(iz,:,:) = -iunit * kZvec(iz) * phi3DFT(iz,:,:)
  end do

  call fft_backward_z(dummy3D)

  avEz = sum( abs(dummy3D)**2, 1 )                 ! save for diagnostics

  call fft_backward_x(dummy3D)
  call fft_backward_y(dummy3D)

  Ez3D = dummy3D



  !-- differentiate potential to find x-component of electric field -- 

  if (jx .eq. 1) then

     avEx = 0.
     Ex3D = 0.

  else

     do ix=1,jx
        dummy3D(:,:,ix) = -iunit * kXvec(ix) * phi3DFT(:,:,ix)
     end do

     call fft_backward_z(dummy3D)

     avEx = sum( abs(dummy3D)**2, 1 )                  ! save for diagnostics

     call fft_backward_x(dummy3D)
     call fft_backward_y(dummy3D)

     Ex3D = dummy3D

  end if


  !-- differentiate potential to find y-component of electric field -- 

  if (jy .eq. 1) then

     avEy = 0.
     Ey3D = 0.

  else

     do iy=1,jy
        dummy3D(:,iy,:) = -iunit * kYvec(iy) * phi3DFT(:,iy,:)
     end do

     call fft_backward_z(dummy3D)

     avEy = sum( abs(dummy3D)**2, 1 )                  ! save for diagnostics

     call fft_backward_x(dummy3D)
     call fft_backward_y(dummy3D)

     Ey3D = dummy3D

  end if


  !-- find phi(center) --


  dummy1Dz(:) = Ez3D(:, jy/2+1, jx/2+1)

  call fft_forward(jz, dummy1Dz, dummy1Dz)
  dummy1Dz = dummy1Dz/float(jz)
  call dealias(jz, dummy1Dz)


  ! amplitude of TOTAL potential at speckle center of sin(kz*z + phase)
  PhiCenter = 2 * abs(dummy1Dz(nz+1) / kZvec(nz+1))

  OmegaBounce = kZvec(nz+1) * sqrt(PhiCenter)
     
  return

end subroutine efield_compute_E
    

!-------------------------------------------------------------------------------------
subroutine efield_evolve_E0()                                   ! input
!-------------------------------------------------------------------------------------

  use parameters, only : phi0, kz, vPhase, dt, iE0Modulation, iWaveFrame, iunit

  implicit none
    
  real                     :: ExpKVdt


  ! in lab frame, E0z has angular frequency vPhase*kz
  if ( iWaveFrame==0 ) then
     ExpKVdt = exp( - iunit*dt*vPhase*kz  )
  else 
     ExpKVdt = 1
  end if
     
  phi0T    = ExpKVdt * phi0T
  E0C(:)   = ExpKVdt * E0C(:)
         
  if (iE0Modulation) then
     E0zC = xMod(jx/2+1)*yMod(jy/2+1)*E0C
  else
     E0zC = E0C
  endif

    
end subroutine efield_evolve_E0

!-------------------------------------------------------------------------------------
subroutine efield_write_spectra()
!-------------------------------------------------------------------------------------

  use parameters, only : runname

  implicit none

  integer    :: ix, iz, ikz, ikx, iky

  real       :: Eksq, phiksq

 !-- densityFTzC(jz) is FT in z at the center (x=0)

  open(12,file = trim(runname)//".EfieldFTz.txt")

  write(12,"('% z-spectra E-field along center line')")
  write(12,"('% 1.kz   2.|E(kz)|^2   3.|Phi(kz)|^2')")
  write(12,*)

  do ikz=2,nkz/2  - nkz/6 -1
     Eksq   = abs(densityFTzC(ikz)/kZvec(ikz)) **2 
     phiksq =  Eksq / kZvec(ikz)**2
     write(12,501) kZvec(ikz), Eksq, phiksq
  end do

  close(12)



  open(12, file = trim(runname)//".EfieldFTx.txt")

  write(12,"('% x-spectra of z-averaged E-field')")
  write(12,"('% 1.kx   2.<|Ex(kx)|^2>   3.<|Ey(kx)|^2>   4.<|Ez(kx)|^2>')")
  write(12,*) 

  iky = 1        !HACK - center slice only, no |k|-averaging

  do ikx=1,nkx/2  - nkx/6 -1

     write(12,501) kxVec(ikx), avEx(iky,ikx), avEy(iky,ikx), avEz(iky,ikx)

  end do

  close(12)

   

501  format(  6400(e12.6,2x)  )

end subroutine efield_write_spectra



!-------------------------------------------------------------------------------------
subroutine efield_get_epsilon(epsilon)                ! output 
!-------------------------------------------------------------------------------------

  use parameters, only : nz, kz, iself

  implicit none

  complex, intent(out)    :: epsilon
  complex                 :: phiInt, phiTot


  phiInt = 2*densityFTzC(1+nz)

  phiTot = phi0T*kz*kz
                 
  if (iself==1) then
     epsilon = phiTot / (phiInt + phiTot)
     !phi = 2*densityFTzC(1+nz)/kz**2 + phi0T
  else
     epsilon = ( phiTot - phiInt ) / phiTot 
     ! = phiExternal/phiTotal = (phiTotal-phiInternal)/phiTotal
  end if


end subroutine efield_get_epsilon

!-------------------------------------------------------------------------------------
subroutine efield_get_E0z_center(E0zCenter)            ! output
!-------------------------------------------------------------------------------------

  real, intent(out)       :: E0zCenter(jz)

  E0zCenter = real(E0zC)

end subroutine efield_get_E0z_center

!-------------------------------------------------------------------------------------

  
end module
